# _*_coding:utf-8_*_
from ATS.instrumentDriver import InstGen
import time

OK = 0
debug_mode = False


class VNA_E5080B(InstGen.VNA):
    def __init__(self, name, info_dict):
        super(VNA_E5080B, self).__init__(name, info_dict)

    def __universalcom(self, cmdlist, func_id=0):
        """
        func_id: 0="write",1="read",2="ask",3="ask for values"
        """
        rval = None
        if not debug_mode:
            for cmd_str in cmdlist:
                if func_id == 0:
                    self.handler.write(cmd_str)
                elif func_id == 1:
                    rval = self.handler.read()
                elif func_id == 2:
                    rval = self.handler.ask(cmd_str)
                elif func_id == 3:
                    rval = self.handler.ask_for_values(cmd_str)
                elif func_id == 4:
                    rval = self.handler.read_values()
        else:
            for cmd in cmdlist:
                if func_id == 0:
                    self.handler.write(cmd)
                    if "?" not in cmd:
                        err = self.handler.ask(":syst:err?")
                        print("###error: ", err)
                elif func_id == 1:
                    rval = self.handler.read()
                elif func_id == 2:
                    rval = self.handler.ask(cmd)
                elif func_id == 3:
                    rval = self.handler.ask_for_values(cmd)
                elif func_id == 4:
                    rval = self.handler.read_values()
                print("\t", rval)
        return rval

    def _init_vna_display(self):
        self.trace_dict = {}
        self.handler.write("TRIG:SOUR MAN")
        trace_id = 1
        self._open_trace("S21", trace_id, "MLOG")
        self.trace_dict["MLOG"] = trace_id
        trace_id += 1
        if "VSWR Variation" in self.rules or "VSWR@Frequency" in self.rules:
            trace_list = []
            self._open_trace("S11", trace_id, "SWR")
            trace_list.append(trace_id)
            trace_id += 1
            self._open_trace("S22", trace_id, "SWR")
            trace_list.append(trace_id)
            trace_id += 1
            self.trace_dict["SWR"] = trace_list

        if "Group Delay Variation" in self.rules or "Group Delay@Frequency" in self.rules:
            self._open_trace("S21", trace_id, "GDEL")
            self.trace_dict["GDEL"] = trace_id
            trace_id += 1

        if "Phase@Frequency" in self.rules:
            self._open_trace("S21", trace_id, "PHAS")
            self.trace_dict["PHAS"] = trace_id
            trace_id += 1

        if "Phase Variation" in self.rules:
            self._open_trace("S21", trace_id, "PHAS")
            self.trace_dict["PHAS_VAR"] = trace_id

    def _recall_sta(self):
        if self.others_set_dict["sta_statu"]:
            self.handler.ask(f"""MMEM:LOAD "{self.others_set_dict['sta_path']}";*opc?""")

    def _clear_errors(self):
        self.handler.write("*CLS")
        self.handler.write("DISP:STAT:LOG:CLE")
        time.sleep(0.3)

    def _open_trace(self, trace_str, windows_id, mes_format):
        if self.handler.ask(f"DISPlay:WINDow{windows_id}:STATe?"):
            self.handler.write(f"DISPlay:WINDow{windows_id}:STATe OFF")
        self.handler.write(f"DISPlay:WINDow{windows_id}:STATe ON")
        self.handler.write(f"CALC{self.channel_id}:PARameter:DEFine 'WINDOWS{windows_id}_{trace_str}',{trace_str},1")
        self.handler.write(f"DISPlay:WINDow{windows_id}:TRACe{windows_id}:FEED 'WINDOWS{windows_id}_{trace_str}'")
        self.handler.write(f"CALC{self.channel_id}:MEAS{windows_id}:FORM {mes_format}")
        t = self.handler.ask("*OPC?")

    def _give_out(self):
        result_dict = {}
        set_index = 0
        self._init_trigger()
        for i in self.rules:
            if i == "VSWR Variation":
                result_dict["VSWR Variation"] = self._mes_var_vswr(self.set_list[set_index]["start_frequency"],
                                                                   self.set_list[set_index]["stop_frequency"],
                                                                   self.trace_dict["SWR"])
            elif i == "Group Delay Variation":
                result_dict["Group Delay Variation"] = self._mes_var_group(self.set_list[set_index]["start_frequency"],
                                                                           self.set_list[set_index]["stop_frequency"],
                                                                           self.trace_dict["GDEL"])
            elif i == "Phase Variation":
                result_dict["Phase Variation"] = self._mes_var_phase(self.set_list[set_index]["start_frequency"],
                                                                     self.set_list[set_index]["stop_frequency"],
                                                                     self.trace_dict["PHAS_VAR"])
            elif i == "VSWR@Frequency":
                result_dict["VSWR@Frequency"] = self._mes_spot_vswr(self.set_list[set_index]["spot_frequency"],
                                                                    self.trace_dict["SWR"])
            elif i == "Group Delay@Frequency":
                result_dict["Group Delay@Frequency"] = self._mes_spot_group(self.set_list[set_index]["spot_frequency"],
                                                                            self.trace_dict["GDEL"])
            elif i == "Phase@Frequency":
                result_dict["Phase@Frequency"] = self._mes_spot_phase(self.set_list[set_index]["spot_frequency"],
                                                                      self.trace_dict["PHAS"])
            elif i == "CF(center)-A":
                result_dict["Bandwith-A"], result_dict["CF(center)-A"], result_dict["fLow-A"], result_dict["fHigh-A"], \
                result_dict["Quality-A"], result_dict["Insertion Loss-A"], result_dict[
                    "Out-Band Rejection-A"] = self._cf_center(self.bw_search_a)
            elif i == "CF(Max)-A":
                result_dict["Bandwith-A"], result_dict["CF(Max)-A"], result_dict["fLow-A"], result_dict["fHigh-A"], \
                result_dict[
                    "Quality-A"], result_dict["Insertion Loss-A"], result_dict["Out-Band Rejection-A"] = self._cf_max(
                    self.bw_search_a)
            elif i == "CF(center)-B":
                result_dict["Bandwith-B"], result_dict["CF(center)-B"], result_dict["fLow-B"], result_dict["fHigh-B"], \
                result_dict["Quality-B"], result_dict["Insertion Loss-B"], result_dict[
                    "Out-Band Rejection-B"] = self._cf_center(self.bw_search_b)
            elif i == "CF(Max)-B":
                result_dict["Bandwith-B"], result_dict["CF(Max)-B"], result_dict["fLow-B"], result_dict["fHigh-B"], \
                result_dict[
                    "Quality-B"], result_dict["Insertion Loss-B"], result_dict["Out-Band Rejection-B"] = self._cf_max(
                    self.bw_search_b)
            set_index += 1
        return result_dict

    def _cf_center(self, bw_search, trace_id=1):
        self.handler.write(f"display:window1:trace{trace_id}:select")
        self._enable_marker_max(1, trace_id)
        bw_search_value = (self.handler.ask(f"CALC{self.channel_id}:MEAS{trace_id}:MARK:BWID:THR?").strip())
        if bw_search_value != bw_search:
            self.handler.ask(f"CALC{self.channel_id}:MEAS{trace_id}:MARK:BWID:THR {bw_search};*opc?")
            self._init_trigger()
        state = int(self.handler.ask(f"CALC{self.channel_id}:MEAS{trace_id}:MARK1:BWID?").strip())

        if not state:
            self.handler.write(f"CALC{self.channel_id}:MEAS{trace_id}:MARK1:BWID ON")
            time.sleep(0.2)
        raw_data = self.handler.ask(f"CALC{self.channel_id}:MEAS{trace_id}:MARK1:BWID:DATA?").strip().split(',')
        loss = round(float(raw_data[-1]), 3)
        bw1 = round(float(raw_data[0]) * 1e-6, 3)  # bandwidth unit = MHz
        f1 = round(float(raw_data[1]) * 1e-6, 3)  # center freq unit = MHz
        quality = round(float(raw_data[2]), 3)
        flow = round(f1 - (bw1 / 2), 3)  # test fLow
        fhigh = round(f1 + (bw1 / 2), 3)  # test fHigh
        # 带外抑制
        self._enable_marker_max(3, trace_id)
        self._enable_marker_max(4, trace_id)
        peak_value = float(self.handler.ask(f"CALC{self.channel_id}:MEAS{trace_id}:MARK3:Y?").strip().split(',')[0])
        self.handler.ask(f"CALC{self.channel_id}:MEAS{trace_id}:MARK3:FUNC:EXEC LPEak;*opc?")
        self.handler.ask(f"CALC{self.channel_id}:MEAS{trace_id}:MARK4:FUNC:EXEC RPEak;*opc?")
        marker3_left_peak = float(
            self.handler.ask(f"CALC{self.channel_id}:MEAS{trace_id}:MARK3:Y?").strip().split(',')[0])
        marker4_right_peak = float(
            self.handler.ask(f"CALC{self.channel_id}:MEAS{trace_id}:MARK4:Y?").strip().split(',')[0])
        if marker4_right_peak > marker3_left_peak and marker4_right_peak != peak_value:
            out_band_rejection = round(peak_value - marker4_right_peak, 3)
        else:
            out_band_rejection = round(peak_value - marker3_left_peak, 3)
        return bw1, f1, flow, fhigh, quality, loss, out_band_rejection

    def _init_trigger_mode(self):
        self.handler.ask("TRIG:SOUR MANual;*OPC?")

    def _init_trigger(self):
        self.handler.ask("INITiate:IMMediate;*OPC?")
        self.handler.ask(f"DISP:WIND:Y:AUTO;*opc?")

    def _cf_max(self, bw_search, trace_id=1):
        self.handler.write(f"display:window1:trace{trace_id}:select")
        self._enable_marker_max(1, trace_id)
        bw_search_value = (self.handler.ask(f"CALC{self.channel_id}:MEAS{trace_id}:MARK:BWID:THR?").strip())
        if bw_search_value != bw_search:
            self.handler.ask(f"CALC{self.channel_id}:MEAS{trace_id}:MARK:BWID:THR {bw_search};*opc?")
            self._init_trigger()
        if not int(self.handler.ask(f"CALC{self.channel_id}:MEAS{trace_id}:MARK1:BWID?")):
            self.handler.write(f"CALC{self.channel_id}:MEAS{trace_id}:MARK1:BWID ON")  # 打开bandWidth search
            time.sleep(0.2)
        raw_data = self.handler.ask(f"CALC{self.channel_id}:MEAS{trace_id}:MARK1:BWID:DATA?").strip().split(',')
        bw1 = round(float(raw_data[0]) * 1e-6, 3)  ## bandwidth unit = MHz
        f1 = round(float(raw_data[1]) * 1e-6, 3)  ## center freq unit = MHz
        quality = round(float(raw_data[2]), 3)
        flow = round(f1 - (bw1 / 2), 3)  # test fLow
        fhigh = round(f1 + (bw1 / 2), 3)  # test fHigh
        marker1_value = self.handler.ask(f"CALC{self.channel_id}:MEAS{trace_id}:MARK1:Y?")
        raw_marker1_loss = marker1_value.split(',')
        loss = round(float(raw_marker1_loss[0]), 3)  # loss
        self._enable_marker_max(2, trace_id)
        raw_cf_max = self.handler.ask(f"CALC{self.channel_id}:MEAS{trace_id}:MARK2:X?").strip().split(
            ',')  # test CF(max)
        cf_max = round(float(raw_cf_max[0]) * 1e-6, 3)
        # 带外抑制
        self._enable_marker_max(3, trace_id)
        self._enable_marker_max(4, trace_id)
        peak_value = float(self.handler.ask(f"CALC{self.channel_id}:MEAS{trace_id}:MARK3:Y?").strip().split(',')[0])
        self.handler.ask(f"CALC{self.channel_id}:MEAS{trace_id}:MARK3:FUNC:EXEC LPEak;*opc?")
        self.handler.ask(f"CALC{self.channel_id}:MEAS{trace_id}:MARK4:FUNC:EXEC RPEak;*opc?")
        marker3_left_peak = float(
            self.handler.ask(f"CALC{self.channel_id}:MEAS{trace_id}:MARK3:Y?").strip().split(',')[0])
        marker4_right_peak = float(
            self.handler.ask(f"CALC{self.channel_id}:MEAS{trace_id}:MARK4:Y?").strip().split(',')[0])
        if marker4_right_peak > marker3_left_peak and marker4_right_peak != peak_value:
            out_band_rejection = round(peak_value - marker4_right_peak, 3)
        else:
            out_band_rejection = round(peak_value - marker3_left_peak, 3)
        return bw1, cf_max, flow, fhigh, quality, loss, out_band_rejection

    def _mes_var_vswr(self, fstart, fstop, trace_list):
        raw_peak1 = self._init_var_mes(trace_list[0], fstart, fstop)
        raw_peak2 = self._init_var_mes(trace_list[1], fstart, fstop)
        peak1 = round(float(raw_peak1[0]), 3)
        peak2 = round(float(raw_peak2[0]), 3)
        return [peak1, peak2]

    def _mes_spot_vswr(self, freq, trace_list):
        raw_data1 = self._init_spot_mes(trace_list[0], freq)
        raw_data2 = self._init_spot_mes(trace_list[1], freq)
        value1 = round(float(raw_data1[0]), 3)
        value2 = round(float(raw_data2[0]), 3)
        return [value1, value2]

    def _mes_var_group(self, fstart, fstop, trace_id=3):
        raw_peak = self._init_var_mes(trace_id, fstart, fstop)
        peak = round(float(raw_peak[0]) * 10 ** 9, 3)
        return peak

    def _mes_spot_group(self, freq, trace_id=3):
        raw_data = self._init_spot_mes(trace_id, freq)
        value = round(float(raw_data[0]) * 10 ** 9, 3)
        return value

    def _mes_var_phase(self, fstart, fstop, trace_id=4):
        raw_peak = self._init_var_mes(trace_id, fstart, fstop)
        peak = round(float(raw_peak[0]), 3)
        return peak

    def _mes_spot_phase(self, freq, trace_id=4):
        raw_data = self._init_spot_mes(trace_id, freq)
        value = round(float(raw_data[0]), 3)
        return value

    def _init_spot_mes(self, trace_id, freq):
        self.handler.write(f"display:window{trace_id}:trace{trace_id}:select")
        self.handler.write(f"CALC{self.channel_id}:MEAS{trace_id}:MARK8 ON")
        self.handler.write(f"CALC{self.channel_id}:MEAS{trace_id}:MARK8:COUP OFF")
        self.handler.write(f"CALC{self.channel_id}:MEAS{trace_id}:MARK8:X {freq}")
        raw_data = self.handler.ask(f"CALC{self.channel_id}:MEAS{trace_id}:MARK8:Y?").strip().split(',')
        return raw_data

    def _init_var_mes(self, trace_id, fstart, fstop):
        self._enable_marker_max(1, trace_id)
        self.handler.write(f"CALC{self.channel_id}:MEAS{trace_id}:MARK:FUNC:DOM:USER:RANG 1")
        self.handler.write(f"CALC{self.channel_id}:MEAS{trace_id}:MARK:FUNC:DOM:USER:START {fstart}")
        self.handler.write(f"CALC{self.channel_id}:MEAS{trace_id}:MARK:FUNC:DOM:USER:STOP {fstop}")
        self.handler.write(f"CALC{self.channel_id}:MEAS{trace_id}:MARK:COUP OFF")
        self.handler.write(f"CALC{self.channel_id}:MEAS{trace_id}:FUNC:STAT ON")
        t = self.handler.ask("*OPC?")
        raw_peak = self.handler.ask(f"CALC{self.channel_id}:MEAS{trace_id}:FUNCtion:DATA?").strip().split(',')
        # peak = round(float(raw_peak[0]), 3)
        return raw_peak

    def _enable_marker_max(self, marker_id, trace_id=1):
        # 查询marker是否开启
        state = int(self.handler.ask(f"CALC{self.channel_id}:MEAS{trace_id}:MARK{marker_id}?").strip())
        if not state:
            self.handler.write(f"CALC{self.channel_id}:MEAS{trace_id}:MARK:COUP OFF")
            self.handler.write(f"CALC{self.channel_id}:MEAS{trace_id}:MARK{marker_id} ON")
        self.handler.ask(f"CALC{self.channel_id}:MEAS{trace_id}:MARK{marker_id}:FUNC:EXEC MAX;*OPC?")
        time.sleep(0.2)
        self.handler.write(f"CALC{self.channel_id}:MEAS{trace_id}:MARK{marker_id}:FUNC:TRAC ON")

    def save_s2p(self, dieId, diePos, folderName, wafer_id):
        # 判断文件是否存在：不存在创建
        folderName = r"{}".format(folderName)
        if not os.path.exists(folderName):
            os.makedirs(folderName)
        filePath = "\{}_{}_{}.s2p".format(wafer_id, dieId, diePos)
        save_path = folderName + filePath
        self.handler.ask(":MMEM:STOR '{}';*opc?".format(save_path))
